# from langchain.embeddings import HuggingFaceBgeEmbeddings
# from langchain_community.embeddings import HuggingFaceBgeEmbeddings
#
#
# class TextToVec():
#     def __init__(self, model_path):
#         self.model = HuggingFaceBgeEmbeddings(
#             model_name=model_path,
#             model_kwargs={'device': 'cpu'},
#             encode_kwargs={'normalize_embeddings': True}  # set True to compute cosine similarity
#         )
#
#     def text2vec(self, text: str):
#         return self.model.embed_query(text)
#
#     def list_text2vec(self, list_text: list):
#         return self.model.embed_documents(list_text)


from sentence_transformers import SentenceTransformer


class TextToVec():
    def __init__(self, model_path):
        # 初始化SentenceTransformer模型
        # model_path可以是本地路径或HuggingFace模型库中的模型名称
        self.model = SentenceTransformer(model_path, device='cpu')

    def text2vec(self, text: str):
        """将单个文本转换为向量"""
        # encode方法默认会返回归一化的向量，适合计算余弦相似度
        return self.model.encode(text)

    def list_text2vec(self, list_text: list):
        """将文本列表转换为向量列表"""
        # 批量处理文本列表，返回对应的向量列表
        return self.model.encode(list_text)
